import cv2
import numpy as np
import os
from pathlib import Path
from tqdm import tqdm

def ensure_dir(path):
    Path(path).mkdir(parents=True, exist_ok=True)

def feather_mask(mask, kernel_size=21):
    """使用高斯模糊羽化mask"""
    # mask = mask.astype(np.float32) / 255.0  # Normalize to [0, 1]
    binary_mask = (mask > 0).astype(np.float32)
    blurred = cv2.GaussianBlur(binary_mask, (kernel_size, kernel_size), 0)
    return np.clip(blurred, 0, 1)

def alpha_blend(component, background, alpha):
    """执行alpha blending"""
    alpha_3ch = np.stack([alpha]*3, axis=-1)
    blended = (component * alpha_3ch + background * (1 - alpha_3ch)).astype(np.uint8)
    return blended

def histogram_equalization(img):
    """对彩色图像执行直方图均衡化"""
    img_yuv = cv2.cvtColor(img, cv2.COLOR_BGR2YUV)
    img_yuv[:, :, 0] = cv2.equalizeHist(img_yuv[:, :, 0])  # 只处理亮度通道
    return cv2.cvtColor(img_yuv, cv2.COLOR_YUV2BGR)

def process_all(image_path, mask_path, feather_path, diagram_path, df_path, fd_path):
    ensure_dir(feather_path)
    ensure_dir(diagram_path)
    ensure_dir(df_path)
    ensure_dir(fd_path)

    image_path = Path(image_path)
    mask_path = Path(mask_path)

    for img_file in tqdm(list(image_path.glob("*"))):
        name = img_file.name
        mask_file = mask_path / name
        if not mask_file.exists():
            print(f"Mask not found for {name}, skipping...")
            continue

        # 读取图像和mask
        img = cv2.imread(str(img_file)).astype(np.float32)
        mask = cv2.imread(str(mask_file), cv2.IMREAD_GRAYSCALE)

        # 羽化 mask
        alpha = feather_mask(mask)

        # -----------------------------
        # 模式1：羽化 + alpha blending
        # 构建 component 图：仅元件区域使用原图，其它为背景图（黑色）
        component = np.zeros_like(img)
        component[mask > 0] = img[mask > 0]
        feather_result = alpha_blend(component, img, alpha)  # 用原图作为背景混合

        cv2.imwrite(str(Path(feather_path) / name), feather_result)

        # -----------------------------
        # 模式2：全图直方图均衡化
        diagram_result = histogram_equalization(img.astype(np.uint8))
        cv2.imwrite(str(Path(diagram_path) / name), diagram_result)

        # -----------------------------
        # 模式3：先均衡化后羽化
        img_eq = histogram_equalization(img.astype(np.uint8)).astype(np.float32)
        df_result = alpha_blend(img_eq, np.full_like(img_eq, 255), alpha)
        cv2.imwrite(str(Path(df_path) / name), df_result)

        # -----------------------------
        # 模式4：先羽化后均衡化
        fd_result = histogram_equalization(feather_result)
        cv2.imwrite(str(Path(fd_path) / name), fd_result)

# 示例用法（请替换为你的路径）

process_all(
    image_path="/home/lhx/pcbsyn/data_zoo/splitsynpcbwithsimplebg/images/test",          # 合成图路径
    mask_path="/home/lhx/pcbsyn/data_zoo/splitsynpcbwithsimplebg/mask/test",            # 元件mask路径
    feather_path="data_zoo/splitsynpcbwithsimplebg/feather",      # 羽化增强结果路径
    diagram_path="data_zoo/splitsynpcbwithsimplebg/diagram",      # 直方图均衡化结果路径
    df_path="data_zoo/splitsynpcbwithsimplebg/df",                # 先均衡再羽化
    fd_path="data_zoo/splitsynpcbwithsimplebg/fd"                 # 先羽化再均衡
)
